//! 池化管理模块
//! 
//! 实现三层池化架构：冷池、温池、热池

use crate::Result;
use crate::types::{SandboxId, PoolType, PoolStats, Language};
use crate::sandbox_manager::Sandbox;
use std::collections::{HashMap, VecDeque};
use std::sync::Arc;
use tokio::sync::RwLock;

/// 池化管理器
pub struct PoolManager {
    /// 冷池 - 存储快照模板
    cold_pool: Arc<RwLock<HashMap<Language, VecDeque<SandboxId>>>>,
    /// 温池 - 预热但未完全初始化的沙箱
    warm_pool: Arc<RwLock<HashMap<Language, VecDeque<Sandbox>>>>,
    /// 热池 - 完全就绪的沙箱
    hot_pool: Arc<RwLock<HashMap<Language, VecDeque<Sandbox>>>>,
    /// 配置
    config: Arc<crate::config::PoolConfig>,
}

impl PoolManager {
    /// 创建池化管理器
    pub fn new(config: Arc<crate::config::PoolConfig>) -> Self {
        Self {
            cold_pool: Arc::new(RwLock::new(HashMap::new())),
            warm_pool: Arc::new(RwLock::new(HashMap::new())),
            hot_pool: Arc::new(RwLock::new(HashMap::new())),
            config,
        }
    }
    
    /// 从热池获取沙箱
    pub async fn acquire_from_hot_pool(&self, language: &Language) -> Result<Option<Sandbox>> {
        let mut hot_pool = self.hot_pool.write().await;
        if let Some(pool) = hot_pool.get_mut(language) {
            if let Some(sandbox) = pool.pop_front() {
                return Ok(Some(sandbox));
            }
        }
        Ok(None)
    }
    
    /// 从温池获取沙箱并激活
    pub async fn acquire_from_warm_pool(&self, language: &Language) -> Result<Option<Sandbox>> {
        let mut warm_pool = self.warm_pool.write().await;
        if let Some(pool) = warm_pool.get_mut(language) {
            if let Some(mut sandbox) = pool.pop_front() {
                // 激活沙箱（标记为运行状态）
                sandbox.metadata.mark_started();
                return Ok(Some(sandbox));
            }
        }
        Ok(None)
    }
    
    /// 释放沙箱到热池
    pub async fn release_to_hot_pool(&self, sandbox: Sandbox) -> Result<()> {
        let language = sandbox.language().clone();
        let mut hot_pool = self.hot_pool.write().await;
        
        // 检查热池大小限制
        let pool = hot_pool.entry(language).or_insert_with(VecDeque::new);
        if pool.len() < self.config.hot_pool_size {
            pool.push_back(sandbox);
        }
        
        Ok(())
    }
    
    /// 获取池统计信息
    pub async fn get_pool_stats(&self, pool_type: PoolType) -> Result<PoolStats> {
        match pool_type {
            PoolType::Hot => {
                let hot_pool = self.hot_pool.read().await;
                let total_count: usize = hot_pool.values().map(|v| v.len()).sum();
                
                Ok(PoolStats {
                    pool_type: PoolType::Hot,
                    total_capacity: self.config.hot_pool_size * hot_pool.len(),
                    used_count: 0, // 热池中的都是空闲的
                    idle_count: total_count,
                    creating_count: 0,
                    error_count: 0,
                    avg_allocation_latency_ms: 5.0, // 热池平均5ms
                    cache_hit_rate: 95.0,
                })
            },
            PoolType::Warm => {
                let warm_pool = self.warm_pool.read().await;
                let total_count: usize = warm_pool.values().map(|v| v.len()).sum();
                
                Ok(PoolStats {
                    pool_type: PoolType::Warm,
                    total_capacity: self.config.warm_pool_size * warm_pool.len(),
                    used_count: 0,
                    idle_count: total_count,
                    creating_count: 0,
                    error_count: 0,
                    avg_allocation_latency_ms: 50.0, // 温池平均50ms
                    cache_hit_rate: 80.0,
                })
            },
            PoolType::Cold => {
                let cold_pool = self.cold_pool.read().await;
                let total_count: usize = cold_pool.values().map(|v| v.len()).sum();
                
                Ok(PoolStats {
                    pool_type: PoolType::Cold,
                    total_capacity: self.config.cold_pool_size * cold_pool.len(),
                    used_count: 0,
                    idle_count: total_count,
                    creating_count: 0,
                    error_count: 0,
                    avg_allocation_latency_ms: 200.0, // 冷池平均200ms
                    cache_hit_rate: 60.0,
                })
            },
        }
    }
    
    /// 预热指定语言的沙箱池
    pub async fn warm_up_pool(&self, language: Language) -> Result<()> {
        // 临时实现 - 后续会实现真正的预热逻辑
        tracing::info!("Warming up pool for language: {}", language);
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::types::{SandboxSpec, Language};
    use crate::config::PoolConfig;
    
    fn create_test_pool_config() -> Arc<PoolConfig> {
        Arc::new(PoolConfig {
            cold_pool_size: 5,
            warm_pool_size: 10,
            hot_pool_size: 20,
            max_pool_size: 100,
            warm_up_concurrency: 5,
            cleanup_interval_secs: 300,
            max_reuse_count: 50,
            default_memory_mb: 0,
            default_cpu_cores: 0,
            default_disk_mb: 0,
            default_timeout_seconds: 0,
            max_idle_seconds: 0,
            max_sandbox_lifetime_seconds: 0,
        })
    }
    
    #[tokio::test]
    async fn test_pool_manager_creation() {
        let config = create_test_pool_config();
        let manager = PoolManager::new(config);
        
        // 测试从空池获取沙箱
        let result = manager.acquire_from_hot_pool(&Language::Python).await.unwrap();
        assert!(result.is_none());
    }
    
    #[tokio::test]
    async fn test_sandbox_release_and_acquire() {
        let config = create_test_pool_config();
        let manager = PoolManager::new(config);
        
        let spec = SandboxSpec {
            language: Language::Python,
            ..Default::default()
        };
        let sandbox = Sandbox::new(spec);
        
        // 释放沙箱到热池
        manager.release_to_hot_pool(sandbox).await.unwrap();
        
        // 从热池获取沙箱
        let acquired = manager.acquire_from_hot_pool(&Language::Python).await.unwrap();
        assert!(acquired.is_some());
        assert_eq!(acquired.unwrap().language(), &Language::Python);
    }
    
    #[tokio::test]
    async fn test_pool_stats() {
        let config = create_test_pool_config();
        let manager = PoolManager::new(config);
        
        let stats = manager.get_pool_stats(PoolType::Hot).await.unwrap();
        assert_eq!(stats.pool_type, PoolType::Hot);
        assert_eq!(stats.avg_allocation_latency_ms, 5.0);
    }
}